package tsj.ai.alibaba.data.analysis.config;

import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.knuddels.jtokkit.api.EncodingType;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Lazy;
import org.springframework.context.annotation.Primary;
import tsj.ai.alibaba.data.analysis.dispatcher.TableRelationDispatcher;
import tsj.ai.alibaba.data.analysis.node.*;
import tsj.ai.alibaba.data.analysis.util.NodeBeanUtil;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicInteger;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static tsj.ai.alibaba.data.analysis.constant.Constant.*;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-13 08:58:43
 */
@Slf4j
@Configuration
@EnableConfigurationProperties({ DataAgentProperties.class})
public class DataAgentConfiguration {

    private static final Logger logger = LoggerFactory.getLogger(DataAgentConfiguration.class);

    @Autowired
    @Lazy
    private KeywordExtractNode keywordExtractNode;

    /**
     * 专用线程池，用于数据库操作的并行处理
     */
    private ExecutorService dbOperationExecutor;

    /**
     * 为了不必要的重复手动配置，不要在此添加其他向量的手动配置，如果扩展其他向量，请阅读spring ai文档
     * <a href="https://springdoc.cn/spring-ai/api/vectordbs.html">...</a>
     * 根据自己想要的向量，在pom文件引入 Boot Starter 依赖即可。此处配置使用内存向量作为兜底配置
     */
    @Primary
    @Bean
    @ConditionalOnMissingBean(VectorStore.class)
    @ConditionalOnProperty(name = "spring.ai.vectorstore.type", havingValue = "simple", matchIfMissing = true)
    public VectorStore simpleVectorStore(EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel).build();
    }


    @Bean
    public StateGraph stateGraph(@Qualifier("dashscopeChatModel") ChatModel chatModel, NodeBeanUtil nodeBeanUtil) throws GraphStateException {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();
            keyStrategyHashMap.put(INPUT_KEY, new ReplaceStrategy());
            keyStrategyHashMap.put(QUERY_REWRITE_NODE_OUTPUT, new ReplaceStrategy());
            return keyStrategyHashMap;
        };
        
        StateGraph stateGraph = new StateGraph(keyStrategyFactory)
                .addNode(QUERY_REWRITE_NODE, nodeBeanUtil.getNodeBeanAsync(QueryRewriteNode.class))
                .addNode(KEYWORD_EXTRACT_NODE, nodeBeanUtil.getNodeBeanAsync(KeywordExtractNode.class))
                .addNode(SCHEMA_RECALL_NODE, nodeBeanUtil.getNodeBeanAsync(SchemaRecallNode.class))
                .addNode(TABLE_RELATION_NODE, nodeBeanUtil.getNodeBeanAsync(TableRelationNode.class))
                .addNode(PLANNER_NODE, nodeBeanUtil.getNodeBeanAsync(PlannerNode.class))

                .addEdge(StateGraph.START, QUERY_REWRITE_NODE)
                .addEdge(QUERY_REWRITE_NODE, KEYWORD_EXTRACT_NODE)
                .addEdge(KEYWORD_EXTRACT_NODE, SCHEMA_RECALL_NODE)
                .addEdge(SCHEMA_RECALL_NODE, TABLE_RELATION_NODE)
                .addConditionalEdges(TABLE_RELATION_NODE, edge_async(new TableRelationDispatcher()),
                        Map.of(PLANNER_NODE, PLANNER_NODE, END, END, TABLE_RELATION_NODE, TABLE_RELATION_NODE)) // retry
                .addEdge(PLANNER_NODE, StateGraph.END);

        // 添加 PlantUML 打印
        GraphRepresentation representation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "analysis flow");
        logger.info("\n=== Data Analysis UML Flow ===");
        logger.info(representation.content());
        logger.info("==================================\n");

        return stateGraph;
    }

    @Bean(name = "dbOperationExecutor")
    public ExecutorService dbOperationExecutor() {
        // 初始化专用线程池，用于数据库操作
        // 线程数量设置为CPU核心数的2倍，但不少于4个，不超过16个
        int threadCount = Math.max(4, Math.min(Runtime.getRuntime().availableProcessors() * 2, 16));
        log.info("Database operation executor initialized with {} threads", threadCount);
        dbOperationExecutor = Executors.newFixedThreadPool(threadCount, new ThreadFactory() {
            private final AtomicInteger threadNumber = new AtomicInteger(1);

            @Override
            public Thread newThread(@NotNull Runnable r) {
                Thread t = new Thread(r, "db-operation-" + threadNumber.getAndIncrement());
                t.setDaemon(true);
                return t;
            }
        });
        return dbOperationExecutor;
    }

    @Bean
    @ConditionalOnMissingBean(BatchingStrategy.class)
    public BatchingStrategy customBatchingStrategy(DataAgentProperties properties) {
        // 使用增强的批处理策略，同时考虑token数量和文本数量限制
        EncodingType encodingType;
        try {
            Optional<EncodingType> encodingTypeOptional = EncodingType
                    .fromName(properties.getEmbeddingBatch().getEncodingType());
            encodingType = encodingTypeOptional.orElse(EncodingType.CL100K_BASE);
        }
        catch (Exception e) {
            log.warn("Unknown encodingType '{}', falling back to CL100K_BASE",
                    properties.getEmbeddingBatch().getEncodingType());
            encodingType = EncodingType.CL100K_BASE;
        }

        return new EnhancedTokenCountBatchingStrategy(encodingType, properties.getEmbeddingBatch().getMaxTokenCount(),
                properties.getEmbeddingBatch().getReservePercentage(),
                properties.getEmbeddingBatch().getMaxTextCount());
    }
}